Skip to content

[RNE Rewrite] Add text and image embeddings pipelines#1292

Open
msluszniak wants to merge 12 commits into
rne-rewritefrom
@ms/add-embeddings
Open

[RNE Rewrite] Add text and image embeddings pipelines#1292
msluszniak wants to merge 12 commits into
rne-rewritefrom
@ms/add-embeddings

Conversation

@msluszniak

@msluszniak msluszniak commented Jun 30, 2026

Copy link
Copy Markdown
Member

Description

Adds text and image embeddings pipelines to the new architecture, achieving parity with the old flow. Embeddings are pure-TypeScript tasks (pooling + L2-norm stay baked into the .pte): text tokenizes and runs forward; image reuses the existing image preprocessor. To run the existing int64-input embedding models unchanged, this adds an int64/Long tensor dtype to the core (the tensor data path is byte-oriented, so it is a small dtype.{h,cpp} + tensor.ts change).

Text inputs are fed at their exact token length (no padding). model.execute validates dynamically-shaped inputs against the [min, max, step] bounds exposed by an optional per-method get_dynamic_dims_<methodName> companion; methods without it keep exact per-dimension validation. This fixes scale-sensitive pooling heads (e.g. DistilUSE's tanh projection), which padding otherwise corrupts.

Includes createTextEmbedder / createImageEmbedder tasks, useTextEmbedder / useImageEmbedder hooks, models.textEmbeddings / models.imageEmbeddings registry entries, an interactive text-embeddings demo in apps/nlp, and a CLIP zero-shot image-embeddings demo in apps/computer-vision.

Introduces a breaking change?

  • Yes
  • No

Type of change

  • Bug fix (change which fixes an issue)
  • New feature (change which adds functionality)
  • Documentation update (improves or adds clarity to existing documentation)
  • Other (chores, tests, code style improvements etc.)

Tested on

  • iOS
  • Android

Testing instructions

  • nlp app → Text Embeddings: seeds a sentence library; type a query and Find similar to rank by cosine similarity, switch models via the chips. Verified on a physical Android device (arm64): all-MiniLM-L6-v2 returns 384-dim L2-normalized embeddings (~25 ms/forward on XNNPACK); DistilUSE ranks correctly with a wide similarity spread (previously compressed by padding).
  • computer-vision app → Image Embeddings: pick an image and rank editable text labels via CLIP zero-shot (image vs. text embeddings). Verified on device.

Screenshots

Related issues

#1247

Checklist

  • I have performed a self-review of my code
  • I have commented my code, particularly in hard-to-understand areas
  • I have updated the documentation accordingly
  • My changes generate no new warnings

Additional notes

All text- and image-embedding models are re-exported with the per-method get_dynamic_dims_<methodName> companion and pinned to v0.10.0: all-MiniLM-L6-v2, all-mpnet-base-v2, multi-qa MiniLM/MPNet, paraphrase-ML, DistilUSE, and CLIP text (plus the DistilUSE/paraphrase CoreML variants). CLIP image has a static input, so it needs no companion and is unchanged.

@msluszniak msluszniak self-assigned this Jun 30, 2026
@msluszniak msluszniak linked an issue Jun 30, 2026 that may be closed by this pull request
@msluszniak msluszniak added the feature PRs that implement a new feature label Jun 30, 2026
Comment thread apps/nlp/app/text-embeddings/index.tsx Outdated
Comment thread apps/nlp/app/text-embeddings/index.tsx Outdated
Comment thread apps/nlp/app/text-embeddings/index.tsx Outdated
Comment thread packages/react-native-executorch/cpp/core/model.cpp Outdated
Comment thread packages/react-native-executorch/cpp/core/model.cpp Outdated
Comment thread apps/nlp/app/text-embeddings/index.tsx
@msluszniak msluszniak marked this pull request as ready for review July 1, 2026 16:07
@msluszniak msluszniak requested a review from barhanc July 1, 2026 16:07
Add int64/Long tensor dtype support and text/image embeddings tasks,
hooks, and model registry entries, plus an interactive text-embeddings
demo screen in apps/nlp.

Closes #1247
model.execute now validates dynamically-shaped forward inputs against the
model-declared [min, max, step] bounds exposed by an optional
get_dynamic_dims method, instead of requiring an exact shape match; models
without it keep exact per-dimension validation. Text embeddings feed the
exact token length with no padding, which fixes scale-sensitive pooling
heads (e.g. DistilUSE's tanh projection).

Point DistilUSE at v0.10.0 (re-exported with get_dynamic_dims).
…mbeddings demo

- Simplify text-embeddings cosine to a dot product (all models L2-normalize)
  and drop redundant inline comments.
- Move the get_dynamic_dims / input-validation contract into the
  ModelHostObject class docs; trim the inline narration in model.cpp.
- Add an Image Embeddings example to the computer-vision app: pick two images
  and compare their CLIP embeddings by cosine similarity.
Rework the computer-vision Image Embeddings screen (based on main's CLIP demo):
pick an image and rank editable text labels by CLIP image/text embedding
similarity, instead of the uninformative two-image score. Pads the scroll
content past the Android nav bar.

Point CLIP text + image at v0.10.0 (text re-exported with get_dynamic_dims;
image unchanged) and declare the textEmbeddings feature in the app.
- model.{h,cpp}: read get_dynamic_dims once per model and cache it instead
  of re-executing the method on every forward() call; reject a present-but-
  malformed declaration (wrong dtype/rank/shape, bad min/max/step, or row
  count not matching forward's tensor input dims) with an explicit error
  instead of silently falling back to exact validation.
- textEmbeddings: throw a clear error when input tokenizes to zero tokens
  (was BigInt(undefined)); fix docstring to match no-padding behavior.
- useTextEmbeddings: expose localPath/tokenizerPath like sibling hooks.
- computer-vision: extract shared skImageToBuffer helper, dedup from
  classification and imageEmbeddings screens.
@msluszniak msluszniak marked this pull request as draft July 3, 2026 14:24
@msluszniak msluszniak force-pushed the @ms/add-embeddings branch from 8e0f200 to ede040e Compare July 3, 2026 14:25
Comment thread packages/react-native-executorch/cpp/core/model.cpp Outdated
Rebase onto rne-rewrite adopted #1296's rewritten model.cpp, which delegates
tensor dtype/shape checks to tensor::fromJs and already supports RangeDim
[min, max, step] bounds. Re-implement variable-length forward inputs on top of
it: parse get_dynamic_dims once per method into cached bounds, build a
SymbolicShape of RangeDims, and pass it to fromJs. Statically shaped methods
keep exact validation.
@msluszniak msluszniak force-pushed the @ms/add-embeddings branch from ede040e to 6c3ccc4 Compare July 3, 2026 14:37
@msluszniak msluszniak marked this pull request as ready for review July 3, 2026 14:37
Align the text/image embeddings tasks with the add-task-pipeline skill (and
every other task): allocate the static output tensor in a `[...] as const`
array, destructure it, and dispose via `tensors.forEach`.

@barhanc barhanc left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Checked only lib implementation. I will take a look at app code on Monday.


Regarding core/ dynamic input support changes, I've added some suggestions that imo make the code more future-proof and easier to change.


Regarding TS side everything is fine, just some small nits. I was thinking though if we don't want to take the opportunity of this refactor and beef-up the TS embeddings side a bit more, something like unifying both image and text embeddings into one pipeline that on top of exposing simple embed method would also expose methods implementing a small vector search data structure like insert, clear, find, etc. I don't know how much effort would that be so it's your call. The current implementation is correct.

Comment thread packages/react-native-executorch/cpp/core/model.cpp
Comment thread packages/react-native-executorch/cpp/core/model.cpp Outdated
Comment thread packages/react-native-executorch/cpp/core/model.cpp Outdated
Comment thread packages/react-native-executorch/src/core/tensor.ts
Comment thread packages/react-native-executorch/cpp/core/model.h Outdated
Comment thread packages/react-native-executorch/src/extensions/nlp/tasks/textEmbeddings.ts Outdated
Comment thread packages/react-native-executorch/src/extensions/nlp/tasks/textEmbeddings.ts Outdated
Comment thread packages/react-native-executorch/src/extensions/nlp/tasks/textEmbeddings.ts Outdated
Comment thread packages/react-native-executorch/src/extensions/nlp/index.ts Outdated
@msluszniak

Copy link
Copy Markdown
Member Author

@barhanc regarding:

  1. Merging text and image pipelines currently I'm slightly towards not doing this. I see the following pros and cons of doing so:

    PROS:

    • Single embed entry and shared result type.
    • Removes a bit of duplication (but only the trivial one).

    CONS:

    • Breaks the split between image embeddings (extensions/cv) and text embeddings (extensions/nlp). Maybe it's incorrect one and both should be under nlp, idk.
    • API gets a bit overloaded, input needs to be string | ImageBuffer union and config will be union of tokenizerPath | preprocessorOpts.
    • text = tokenizer + int64 ids/mask + truncation; image = image preprocessor + float NCHW.
    • Even CLIP has text and image part and each is a separate *.pte file.
  2. Regarding small vector DB functionality, I would do it in a separate PR, but not as a part of the embedding flow, but as a small, separate functionality. That is because otherwise it would introduce stateful functionality inside the generally stateless pipelines and couples them. It is also a separate scope from embeddings themselves.

If you think that there are other things to consider regarding these two, please let me know and I will re-think it. In a meantime, I will address inline comments :).

Address review: parse dynamic input shapes into a SymbolicShape-per-input
(1:1 with inputs, empty for non-tensor) via parseDynamicInputShapes, keyed on a
per-method companion get_dynamic_dims_<methodName> (int32 [rank,3]). Build the
methodName -> shapes map for all methods at construction instead of hardcoding
forward; execute just indexes by input position. unwrap now throws
std::runtime_error (no Runtime needed).
Address review:
- createTextEmbeddings -> createTextEmbedder, createImageEmbeddings ->
  createImageEmbedder; files textEmbedding.ts/imageEmbedding.ts (no 's'),
  hooks useTextEmbedder/useImageEmbedder, and *EmbedderModel types
- forward/forwardWorklet -> embed/embedWorklet
- t-prefix tensor vars (tTokenIds/tAttentionMask); document input truncation
- stop exporting tasks from extensions/<domain>/index.ts (drop textEmbedding
  and the pre-existing tokenization export)
Comment thread packages/react-native-executorch/cpp/core/model.cpp
Comment thread packages/react-native-executorch/cpp/core/model.cpp Outdated
Comment thread packages/react-native-executorch/cpp/core/model.cpp
Comment thread packages/react-native-executorch/src/hooks/useTextEmbedder.ts Outdated

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why was this change only added to classification when other CV screens use the same pattern?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It was an unrelated cleanup that I'd only applied here — reverted it to keep this PR scoped to embeddings.

@barhanc

barhanc commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Tested apps and they work fine. Also I see that on huggingface we also have LFM text embeddings, why are they not added here?

@barhanc

barhanc commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

One more thing, we should probably add some information about the dynamic methods to skills and jsdoc so that agents are aware of the companion method requirement, this is adhoc solution until we have a proper documentation.

- Add a jsi::JSError overload of unwrap and use it across the JSI host
  functions (getMethodNames/getMethodMeta/execute), so runtime failures
  surface as JS errors; the std::runtime_error overload stays for the
  load-time constructor path where no jsi::Runtime is available.
- Validate output placeholders against the tensor's actual produced shape
  (output.toTensor().sizes()) so dynamically-shaped outputs are supported.
- useTextEmbedder: report the model download progress directly instead of
  averaging with the tiny tokenizer download.
- Revert the unrelated classification screen refactor (skImageToBuffer /
  insets) to keep this PR scoped to embeddings.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature PRs that implement a new feature refactoring

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[RNE Rewrite] Add image and text embeddings pipelines

2 participants